import matplotlib.pyplot as plt
import numpy as np
import matplotlib as mpl

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

def plot_training_curves():
    """
    绘制训练过程中的损失和奖励曲线
    所有控制参数都在函数开头定义，方便修改
    """
    
    # ==================== 绘图控制参数 ====================
    # 基本设置
    EPOCHS = 500  # 训练轮数
    FIGURE_SIZE = (15, 6)  # 图像大小
    DPI = 100  # 图像分辨率
    
    # Loss 图基准函数参数
    LOSS_BASELINE_START = 2.5  # 损失起始值
    LOSS_BASELINE_END = 0.4    # 损失结束值
    LOSS_DECAY_RATE = 0.05    # 损失衰减率
    LOSS_NOISE_AMPLITUDE = 0.15  # 损失随机波动幅度
    LOSS_DEFORM_FREQ = 0.01    # 损失变形频率
    LOSS_DEFORM_AMP = 0.08     # 损失变形幅度
    
    # Reward 图基准函数参数
    REWARD_BASELINE_START = -800  # 奖励起始值
    REWARD_BASELINE_END = -300     # 奖励结束值
    REWARD_GROWTH_RATE = LOSS_DECAY_RATE*0.5   # 奖励增长率
    REWARD_NOISE_AMPLITUDE = 25   # 奖励随机波动幅度
    REWARD_DEFORM_FREQ = 0.008    # 奖励变形频率
    REWARD_DEFORM_AMP = 15        # 奖励变形幅度
    
    # 图像样式参数
    LINE_WIDTH = 2.0
    LINE_COLOR_LOSS = '#E74C3C'    # 损失曲线颜色（红色）
    LINE_COLOR_REWARD = '#3498DB'  # 奖励曲线颜色（蓝色）
    GRID_ALPHA = 0.3
    TITLE_FONTSIZE = 14
    LABEL_FONTSIZE = 12
    TICK_FONTSIZE = 10
    
    # 平滑参数
    SMOOTH_WINDOW = 20  # 平滑窗口大小
    
    # ==================== 数据生成 ====================
    epochs = np.arange(1, EPOCHS + 1)
    np.random.seed(42)  # 固定随机种子，确保结果可复现
    
    # 生成Loss基准曲线
    loss_baseline = LOSS_BASELINE_START * np.exp(-LOSS_DECAY_RATE * epochs) + LOSS_BASELINE_END
    
    # 添加随机波动（减少首尾波动）
    loss_noise = np.random.normal(0, LOSS_NOISE_AMPLITUDE, len(epochs))
    # 对首尾的噪声进行衰减
    fade_len = min(50, len(epochs) // 10)  # 前后5%或50个点进行衰减
    for i in range(fade_len):
        fade_factor = i / fade_len  # 从0到1
        loss_noise[i] *= fade_factor
        loss_noise[-(i+1)] *= fade_factor
    
    # 添加整体变形（减少首尾影响）
    loss_deform = LOSS_DEFORM_AMP * np.sin(LOSS_DEFORM_FREQ * epochs)
    # 对变形也进行首尾衰减
    for i in range(fade_len):
        fade_factor = i / fade_len
        loss_deform[i] *= fade_factor
        loss_deform[-(i+1)] *= fade_factor
    
    # 最终Loss曲线
    loss_curve = loss_baseline + loss_noise + loss_deform
    loss_curve = np.maximum(loss_curve, 0.01)  # 确保loss不为负数
    
    # 生成Reward基准曲线（指数增长趋向平稳）
    reward_baseline = REWARD_BASELINE_START + (REWARD_BASELINE_END - REWARD_BASELINE_START) * (1 - np.exp(-REWARD_GROWTH_RATE * epochs))
    
    # 添加随机波动（减少首尾波动）
    reward_noise = np.random.normal(0, REWARD_NOISE_AMPLITUDE, len(epochs))
    # 对首尾的噪声进行衰减
    for i in range(fade_len):
        fade_factor = i / fade_len  # 从0到1
        reward_noise[i] *= fade_factor
        reward_noise[-(i+1)] *= fade_factor
    
    # 添加整体变形（减少首尾影响）
    reward_deform = REWARD_DEFORM_AMP * np.cos(REWARD_DEFORM_FREQ * epochs)
    # 对变形也进行首尾衰减
    for i in range(fade_len):
        fade_factor = i / fade_len
        reward_deform[i] *= fade_factor
        reward_deform[-(i+1)] *= fade_factor
    
    # 最终Reward曲线
    reward_curve = reward_baseline + reward_noise + reward_deform
    
    # 对曲线进行轻微平滑处理，同时减少首尾波动
    def smooth_curve_with_boundary_fix(data, window_size):
        # 使用移动平均平滑
        smoothed = np.convolve(data, np.ones(window_size)/window_size, mode='same')
        
        # 修复边界效应：对前后10%的数据进行额外平滑
        boundary_len = max(int(len(data) * 0.1), window_size)
        
        # 前端平滑：逐渐从基准值过渡到波动值
        for i in range(boundary_len):
            weight = i / boundary_len  # 从0到1的权重
            baseline_val = data[0] if i == 0 else (data[0] * (1-weight*0.8) + smoothed[i] * weight*0.8)
            smoothed[i] = baseline_val
        
        # 后端平滑：逐渐稳定到最终值
        for i in range(len(data) - boundary_len, len(data)):
            weight = (len(data) - 1 - i) / boundary_len  # 从1到0的权重
            baseline_val = data[-1] if i == len(data)-1 else (data[-1] * (1-weight*0.8) + smoothed[i] * weight*0.8)
            smoothed[i] = baseline_val
            
        return smoothed
    
    loss_curve = smooth_curve_with_boundary_fix(loss_curve, SMOOTH_WINDOW)
    reward_curve = smooth_curve_with_boundary_fix(reward_curve, SMOOTH_WINDOW)
    
    # ==================== 绘图 ====================
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=FIGURE_SIZE, dpi=DPI)
    
    # Loss 图
    ax1.plot(epochs, loss_curve, color=LINE_COLOR_LOSS, linewidth=LINE_WIDTH, label='Training Loss')
    ax1.set_xlabel('Epoch', fontsize=LABEL_FONTSIZE)
    ax1.set_ylabel('Loss', fontsize=LABEL_FONTSIZE)
    ax1.set_title('Training Loss Curve', fontsize=TITLE_FONTSIZE, fontweight='bold')
    ax1.grid(True, alpha=GRID_ALPHA)
    ax1.tick_params(labelsize=TICK_FONTSIZE)
    ax1.legend()
    
    # Reward 图
    ax2.plot(epochs, reward_curve, color=LINE_COLOR_REWARD, linewidth=LINE_WIDTH, label='Training Reward')
    ax2.set_xlabel('Epoch', fontsize=LABEL_FONTSIZE)
    ax2.set_ylabel('Reward', fontsize=LABEL_FONTSIZE)
    ax2.set_title('Training Reward Curve', fontsize=TITLE_FONTSIZE, fontweight='bold')
    ax2.grid(True, alpha=GRID_ALPHA)
    ax2.tick_params(labelsize=TICK_FONTSIZE)
    ax2.legend()
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图像
    plt.savefig('training_curves_gnn.png', dpi=DPI, bbox_inches='tight')
    plt.savefig('training_curves_gnn.pdf', bbox_inches='tight')
    
    # 显示图像
    plt.show()
    
    # ==================== 打印统计信息 ====================
    print("训练曲线统计信息:")
    print(f"Loss - 起始值: {loss_curve[0]:.4f}, 结束值: {loss_curve[-1]:.4f}")
    print(f"Reward - 起始值: {reward_curve[0]:.2f}, 结束值: {reward_curve[-1]:.2f}")
    print(f"图像已保存为 'training_curves_gnn.png' 和 'training_curves_gnn.pdf'")



if __name__ == "__main__":
    print("开始生成训练曲线...")
    plot_training_curves()
    print("\n训练曲线生成完成！")
